package com.zt.bigdata.elasticsearch.utils.transport.apis;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.zt.bigdata.elasticsearch.utils.transport.ESClient;
import org.apache.commons.collections.CollectionUtils;
import org.elasticsearch.action.admin.cluster.storedscripts.DeleteStoredScriptResponse;
import org.elasticsearch.action.admin.cluster.storedscripts.GetStoredScriptResponse;
import org.elasticsearch.action.admin.cluster.storedscripts.PutStoredScriptResponse;
import org.elasticsearch.action.index.IndexRequestBuilder;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchRequestBuilder;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.transport.TransportClient;
import org.elasticsearch.common.bytes.BytesArray;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.script.ScriptType;
import org.elasticsearch.script.mustache.SearchTemplateRequestBuilder;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.aggregations.AggregationBuilder;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.sort.FieldSortBuilder;
import org.elasticsearch.search.sort.SortOrder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.data.domain.Pageable;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * @ClassName SearchApi
 * @Description
 * @Author zhangtonghy
 * @Date 2019-02-21 16:17
 * @Copyright: 版权所有 (C) zt zt.
 * @注意：本内容仅限于zt内部传阅，禁止外泄以及用于其他的商业目的
 **/
public class SearchApi {
    private TransportClient transportClient;
    private ObjectMapper mapper = new ObjectMapper();
    private static final Logger logger = LoggerFactory.getLogger(SearchApi.class);

    private TransportClient getClient() {
        if (null == transportClient) {
            ESClient esClient = new ESClient("cdn-bigdata", "39.134.86.164:9300");
            transportClient = esClient.getClient();
        }
        return transportClient;
    }

    public SearchResponse searchByQuery(String index, QueryBuilder queryBuilder, QueryBuilder filter, Pageable pageable) {
        SearchRequestBuilder searchRequestBuilder = getClient().prepareSearch(index)
                .setQuery(queryBuilder)
                .setPostFilter(filter)
                .setExplain(true);
        if (null != pageable) {
            searchRequestBuilder
                    .setFrom((pageable.getPageNumber() - 1) * pageable.getPageSize())
                    .setSize(pageable.getPageSize());
        }
        return searchRequestBuilder.get();
    }

    public List<Map<String, Object>> scrollSearchByQuery(String index, QueryBuilder queryBuilder) {
        List<Map<String, Object>> resp = new ArrayList<>();
        SearchResponse scrollResp = getClient()
                .prepareSearch(index)
                .addSort(FieldSortBuilder.DOC_FIELD_NAME, SortOrder.ASC)
                .setScroll(new TimeValue(600000))
                .setQuery(queryBuilder)
                .setSize(10000)
                .get();
        //Scroll until no hits are returned
        do {
            for (SearchHit hit : scrollResp.getHits().getHits()) {
                resp.add(hit.getSourceAsMap());
            }
            scrollResp = getClient()
                    .prepareSearchScroll(scrollResp.getScrollId())
                    .setScroll(new TimeValue(600000))
                    .execute()
                    .actionGet();

        } while (scrollResp.getHits().getHits().length != 0);
        return resp;
    }

    public SearchResponse searchAndAgg(String index, QueryBuilder queryBuilder, List<AggregationBuilder> aggregationBuilders) {

        SearchRequestBuilder searchRequestBuilder = getClient().prepareSearch(index)
                .setQuery(queryBuilder);

        if (CollectionUtils.isNotEmpty(aggregationBuilders)) {
            aggregationBuilders.forEach(aggBuilder -> searchRequestBuilder.addAggregation(aggBuilder));
        }
        return searchRequestBuilder.get();
    }


    public AggregationBuilder generateAgg(String aggName, String field) {
        return AggregationBuilders.dateHistogram(aggName)
                .field(field);
    }


    /**
     * @param templateId
     * @param content    "{ \"query\": { \"match\" : { \"gender\" : \"{{param_gender}}\" } }}"
     * @return
     */
    public PutStoredScriptResponse createSearchTemplate(String templateId, String content) {


        PutStoredScriptResponse response = getClient()
                .admin()
                .cluster()
                .preparePutStoredScript()
                .setId(templateId)
                .setContent(new BytesArray(content), XContentType.JSON)
                .get();

        return response;
    }

    public GetStoredScriptResponse getSearchTemplate(String templateId) {
        return getClient()
                .admin()
                .cluster()
                .prepareGetStoredScript()
                .setId(templateId)
                .get();
    }

    public DeleteStoredScriptResponse deleteSearchTemplate(String templateId) {
        return getClient().admin().cluster().prepareDeleteStoredScript()
                .setId(templateId)
                .get();
    }

    public SearchResponse searchByTemplate(Map<String, Object> templateParams, String templateId) {

        SearchResponse sr = new SearchTemplateRequestBuilder(getClient())
                .setScript(templateId)
                .setScriptType(ScriptType.STORED)
                .setScriptParams(templateParams)
                .setRequest(new SearchRequest())
                .get()
                .getResponse();
        return sr;
    }

    /**
     * @param templateParams
     * @param script         "{ \"query\": { \"match\" : { \"gender\" : \"{{param_gender}}\" } }}"
     * @return
     */
    public SearchResponse searchByInlineTemplate(Map<String, Object> templateParams, String script) {

        SearchResponse sr = new SearchTemplateRequestBuilder(getClient())
                .setScript(script)
                .setScriptType(ScriptType.INLINE)
                .setScriptParams(templateParams)
                .setRequest(new SearchRequest())
                .get()
                .getResponse();
        return sr;
    }

    public static void main(String[] args) {


//        SearchApi searchApi = new SearchApi();
//
//        Map<String, Object> map = new HashMap();
//        map.put("param_gender", "M");
//        SearchResponse searchResponse = searchApi
//                .searchByInlineTemplate(map, "{ \"query\": { \"match\" : { \"gender\" : \"{{param_gender}}\" } }}");
//        for (int i = 0; i < searchResponse.getHits().getHits().length; i++) {
//            System.out.println(searchResponse.getHits().getHits()[i].getSourceAsMap());
//        }

        String index = "hy-bigdata-2019021800";

        SearchApi searchApi = new SearchApi();
        searchApi.getClient().prepareClearScroll();

        DocApi docApi = new DocApi();

        SearchResponse scrollResp = searchApi.getClient()
                .prepareSearch(index)
                .setScroll(new TimeValue(600000))
                .setQuery(QueryBuilders.matchAllQuery())
                .setSize(10000)
                .get();
        //Scroll until no hits are returned
        do {
            List<IndexRequestBuilder> resp = new ArrayList<>();

            for (SearchHit hit : scrollResp.getHits().getHits()) {
                IndexRequestBuilder indexRequestBuilder = searchApi.getClient().prepareIndex(index, "log").setSource(hit.getSourceAsMap());
                resp.add(indexRequestBuilder);
            }
            docApi.bulkRequest(resp);
            System.out.println("插入"+resp.size());
            scrollResp = searchApi.getClient()
                    .prepareSearchScroll(scrollResp.getScrollId())
                    .setScroll(new TimeValue(600000))
                    .execute()
                    .actionGet();

        } while (scrollResp.getHits().getHits().length != 0);

        searchApi.getClient().prepareClearScroll();



    }
}
